package main

import (
	"bufio"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/hex"
	"flag"
	"fmt"
	"hash"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"

	_ "embed"

	"gitee.com/MM-Q/colorlib"
)

// 定义一个结构体, 用于存储Help文件的信息
//
//go:embed help.txt
var help string

// 定义一个字符串变量，用于存储版本信息
var version = "v1.0.4"

// 定义一个字符串切片变量，用于存储命令行参数
var args []string

// 定义一个颜色库变量, 用于输出带颜色的文本
var color *colorlib.ColorLib = colorlib.NewColorLib()

// 创建一个结构体用于在通道中传递文件信息
type hashInfo struct {
	hash string
	path string
}

// 创建一个存储源基准校验码文件的Map
var actualMap sync.Map

// 定义命令行参数变量
var (
	// 校验文件 - 用于检查文件的完整性 优先级最高
	checkF = flag.Bool("c", false, "根据校验文件检查文件完整性")
	// 用于计算文件或目录的大小
	calcSizeF = flag.Bool("cs", false, "计算文件或目录的大小(单位: B, KB, MB, GB, TB)")

	// 生成校验文件 - 用于生成文件的校验码
	recurseF = flag.Bool("r", false, "开启递归扫描目录")
	jobsF    = flag.Int("j", 3, "设置并发任务数")
	writeF   = flag.Bool("w", false, "将校验结果保存到文件中")
	outPutF  = flag.String("o", "checksum.txt", "指定校验文件的保存路径")
	sizeF    = flag.Bool("size", false, "在校验或生成哈希值时显示文件大小(单位: B, KB, MB, GB, TB)")

	// 校验文件的哈希算法
	md5F    = flag.Bool("md5", false, "使用MD5算法校验文件")
	sha1F   = flag.Bool("sha1", false, "使用SHA1算法校验文件")
	sha256F = flag.Bool("sha256", false, "使用SHA256算法校验文件")

	// 显示版本信息
	versionF = flag.Bool("v", false, "显示版本信息")
)

func main() {
	// 初始化程序
	if err := mainInit(); err != nil {
		color.Red(err.Error()) // 打印错误信息
		flag.PrintDefaults()   // 打印帮助信息
		os.Exit(1)             // 退出程序
	}

	// 如果使用了 -v 参数，则显示版本信息并退出程序
	if *versionF {
		color.Greenf("fastcheck %s", version)
		return
	}

	// 检查是否有命令行参数，如果没有则退出程序
	if len(args) == 0 {
		color.Red("error: 请输入参数选项或文件路径，或者使用 -h 查看帮助信息")
		flag.PrintDefaults()
		return
	}

	// 如果使用了 -cs 参数，则计算文件或目录的大小
	if *calcSizeF {
		// 调用主函数计算大小
		if err := calculateSize(args, *jobsF); err != nil {
			color.Redf("error: %v\n", err)
			return
		}
		return
	}

	/* 以下为校验数据的模块 */
	// 根据命令行参数选择哈希算法并调用 mainCheck
	hashFunc := selectHashFunction()

	// 根据命令行参数选择校验或生成校验码的运行方式
	if err := mainCheck(hashFunc); err != nil {
		color.Red(err.Error())
		return
	}
}

// mainCheck 根据命令行参数选择校验或生成校验码的运行方式
func mainCheck(hashFunc func() hash.Hash) error {
	// 如果使用了 -c 参数，则检查文件完整性
	if *checkF {
		if err := checkIntegrity(hashFunc); err != nil {
			return err
		}
	}

	// 正常校验模式
	if err := generateChecksums(hashFunc, args); err != nil {
		return err
	}

	return nil
}

// selectHashFunction 根据命令行参数选择哈希算法
func selectHashFunction() func() hash.Hash {
	if *md5F { // 如果使用了 -md5 参数
		return md5.New
	} else if *sha1F { // 如果使用了 -sha1 参数
		return sha1.New
	} else if *sha256F { // 如果使用了 -sha256 参数
		return sha256.New
	}
	return md5.New // 默认使用 MD5 算法
}

func mainInit() error {
	// 定义命令行参数
	flag.Usage = func() {
		fmt.Println(help) // 打印帮助信息
	}

	// 初始化解析命令行参数
	flag.Parse()

	// 获取所有命令行参数
	args = flag.Args()

	// 遍历并规范化路径
	for i, path := range args {
		args[i] = filepath.Clean(path) // 路径规范化解析
	}

	// 检查命令行参数
	if err := checkHashFlags(); err != nil {
		return err
	}

	// 检查并发任务数
	if *jobsF <= 1 {
		return fmt.Errorf("error: 并发任务数必须大于 1")
	}

	return nil
}

// 检查命令行hash算法参数
func checkHashFlags() error {
	count := 0
	// 检查是否有哈希算法参数
	if *md5F {
		count++
	}

	// 检查是否有哈希算法参数
	if *sha1F {
		count++
	}

	// 检查是否有哈希算法参数
	if *sha256F {
		count++
	}

	// 如果没有哈希算法参数，则默认使用 MD5 算法
	if count == 0 {
		*md5F = true
	}

	// 如果有多个哈希算法参数，则退出程序
	if count > 1 {
		flag.PrintDefaults()
		return fmt.Errorf("error: 只能使用一个哈希算法")
	}

	return nil
}

// 通用函数，检查文件或目录是否存在
func checkFileOrDir(path string) (os.FileInfo, error) {
	// 路径规范化解析
	path = filepath.Clean(path)

	// 获取文件或目录的信息
	fileInfo, err := os.Stat(path)
	if err != nil {
		// 如果文件或目录不存在，则返回错误
		if os.IsNotExist(err) {
			return nil, fmt.Errorf("error: 文件或目录 %s 不存在", path)
			// 如果文件或目录不可访问，则返回错误
		} else if os.IsPermission(err) {
			return nil, fmt.Errorf("error: 无权限访问 %s", path)
		}
		// 如果其他错误，则返回错误
		return nil, fmt.Errorf("error: 无法访问 %s: %v", path, err)
	}
	return fileInfo, nil
}

// generateChecksums 检查文件或目录的 MD5 值
func generateChecksums(hashFunc func() hash.Hash, filePaths []string) error {
	var (
		outPutFile *os.File       // 存储文件句柄
		err        error          // 存储错误信息
		wg         sync.WaitGroup // 等待组
		chMap      chan hashInfo  // 哈希通道
	)

	// 扩展文件路径，支持通配符和正则表达式
	expandedFilePaths, err := expandFilePaths(filePaths)
	if err != nil {
		return fmt.Errorf("error: 扩展文件路径失败: %v", err)
	}

	// 创建一个通道用于限制并发数
	job := make(chan string, *jobsF) // 限制并发数 默认3

	// 如果是校验模式则初始化一个通道用于接收校验结果
	if *checkF {
		chMap = make(chan hashInfo, *jobsF) // 限制并发数 默认3
	}

	// 如果使用了 -w 参数，则创建文件
	if *writeF {
		// 创建保存结果的文件
		outPutFile, err = os.Create(*outPutF)
		if err != nil {
			color.Redf("error: 无法创建文件 %s: %v\n", "md5sum.txt", err)
			return err
		}
		defer outPutFile.Close()
	}

	// 启动固定数量的工作协程
	for i := 0; i < cap(job); i++ { // 根据通道的容量限制协程数量
		wg.Add(1)
		go func() {
			defer wg.Done()
			for path := range job {
				// 计算文件的 MD5 值
				hash, err := checksum(path, hashFunc)
				if err != nil {
					color.Redf("error: 计算文件 %s 的哈希值失败: %v", path, err)
					continue
				}
				// 根据命令行参数判断是否输出到控制台或文件
				if *writeF {
					if *checkF {
						// 如果是校验模式，则将结果保存到哈希通道中
						chMap <- hashInfo{hash: hash, path: path}
					} else {
						msg := fmt.Sprintf("%s\t%q\n", hash, path) // 拼接哈希值和文件路径
						// 如果是校验的是存储校验码的文件，则跳过
						if strings.Contains(path, *outPutF) {
							continue
						}

						// 如果是普通模式，则将结果写入文件
						if _, err := outPutFile.WriteString(msg); err != nil {
							color.Redf("error: 写入文件时出错: %v\n", err)
							continue
						}
					}
				} else {
					if *checkF {
						// 如果是校验模式，则将结果保存到哈希通道中
						chMap <- hashInfo{hash: hash, path: path}
					} else {
						// 如果是普通模式，则直接输出到控制台
						if *sizeF {
							size, err := getSize(path) // 获取文件大小
							if err != nil {            // 如果获取文件大小失败，则输出错误信息
								fmt.Printf("%-10s\t%-10s\t%-20s\n", hash, err.Error(), path)
								continue
							}
							sizeNew := humanReadableSize(size) // 格式化文件大小
							fmt.Printf("%-10s\t%-10s\t%-20s\n", hash, sizeNew, path)
						} else {
							fmt.Printf("%-10s\t\t%-20s\n", hash, path)
						}
					}
				}
			}
		}()
	}

	// 声明一个等待组
	var wh sync.WaitGroup
	// 启动一个协程, 用于写入校验结果到actualMap中
	if *checkF {
		wh.Add(1)
		go func() {
			defer wh.Done()
			for hashInfo := range chMap { // 从通道中读取hashInfo结构体
				// 将结果保存到actualMap中 key为文件路径 value为哈希值
				actualMap.Store(hashInfo.path, hashInfo.hash)
			}
		}()
	}

	// 标记通道是否已经关闭
	closed := false

	// 遍历所有文件路径
	for _, path := range expandedFilePaths {
		fileInfo, err := checkFileOrDir(path)
		if err != nil {
			color.Red(err.Error())
			continue
		}

		if fileInfo.IsDir() {
			if *recurseF {
				err := filepath.WalkDir(path, func(subPath string, d os.DirEntry, err error) error {
					// 规范化路径
					subPath = filepath.Clean(subPath)
					if err != nil {
						color.Redf("error: 无法访问 %s: %v", subPath, err)
						return err
					}
					if !d.IsDir() {
						job <- subPath
					}
					return nil
				})
				if err != nil {
					color.Redf("error: 递归扫描目录 %s 时出错: %v", path, err)
					if !closed {
						close(job) // 提前关闭通道
					}
					wg.Wait() // 等待所有协程完成
					return err
				}
			} else {
				// 如果不是递归模式，则输出警告信息
				color.Yellowf("Warn: %s 是一个目录，非文件。请使用 -r 或 --recurse 开启递归扫描目录", path)
			}
		} else {
			// 如果是文件，则将文件路径发送到通道
			job <- path
		}
	}

	// 等待所有任务加入队列后关闭通道
	if !closed {
		close(job) // 提前关闭通道
	}
	// 等待所有工作协程完成
	wg.Wait()

	// 等待所有写入到actualMap的协程完成
	if *checkF {
		close(chMap) // 在所有写入操作完成后关闭通道
		wh.Wait()    // 等待写入 actualMap 的协程完成
	}

	// 根据命令行参数判断是否输出到控制台或文件
	if *writeF {
		color.Greenf("success: 生成校验码成功，保存到文件 %s", *outPutF)
	}

	// 如果是普通模式，返回空哈希表
	return nil
}

// 校验模式，根据校验文件检查文件完整性
func checkIntegrity(hashFunc func() hash.Hash) error {
	// 获取校验文件路径
	checkFile := args[0]

	// 用于存储文件路径的切片
	filePaths := []string{}

	// 打开校验文件
	file, err := os.Open(checkFile)
	if err != nil {
		return fmt.Errorf("error: 打开校验文件 %s 失败: %v", checkFile, err)
	}
	defer file.Close()

	// 读取校验文件，存储到 map 中
	expectedMap := make(map[string]string)
	scanner := bufio.NewScanner(file) // 创建一个扫描器

	// 定义一个计数器
	var lineCount int

	// 逐行读取校验文件
	for scanner.Scan() { // 逐行读取
		line := scanner.Text() // 获取当前行的文本

		lineCount++ // 计数器加 1

		if line == "" { // 如果当前行是空行，则跳过
			continue
		}

		// 解析校验文件中的哈希值和文件路径
		parts := strings.Fields(line) // 按空格分割
		if len(parts) != 2 {          // 如果分割后的长度不为 2，则跳过
			color.Redf("error: 校验文件格式错误, 文件 %s 的第 %d 行, %s", checkFile, lineCount, line)
			continue
		}
		expectedHash := parts[0]                 // 哈希值
		filePath := strings.Join(parts[1:], " ") // 文件路径

		// 去除路径中的引号
		filePath = strings.Trim(filePath, `"`)

		// 存储到 map 中
		expectedMap[filePath] = expectedHash

		// 存储文件路径到切片中
		filePaths = append(filePaths, filePath)
	}

	if err := scanner.Err(); err != nil {
		return fmt.Errorf("error: 读取校验文件时出错: %v", err)
	}

	// 调用 generateChecksums 函数计算哈希值
	err = generateChecksums(hashFunc, filePaths)
	if err != nil {
		return fmt.Errorf("error: 计算哈希值失败: %v", err)
	}

	// 对比哈希值
	for filePath, expectedHash := range expectedMap {
		// 获取实际的哈希值
		actualHash, ok := actualMap.Load(filePath)
		if !ok {
			// 如果实际的哈希值不存在，则跳过
			continue
		}

		if actualHash == expectedHash { // 比较哈希值
			color.Greenf("%s: OK", filePath)
		} else {
			color.Redf("%s: FAILED (预期Hash值: %s, 实际Hash值: %s)", filePath, expectedHash, actualHash)
		}
	}

	return nil
}

// 核心函数 checksum 计算文件的哈希值，支持多种算法
func checksum(filepath string, hashFunc func() hash.Hash) (string, error) {
	// 打开文件
	file, err := os.Open(filepath)
	if err != nil {
		color.Redf("error: 打开文件 %s 失败: %v", filepath, err)
		return "", fmt.Errorf("error: 打开文件 %s 失败: %v", filepath, err)
	}
	defer file.Close()

	// 创建哈希对象
	hash := hashFunc()

	// 手动分块读取文件内容
	buffer := make([]byte, 4*1024*1024) // 4MB 的缓冲区
	for {
		n, err := file.Read(buffer)
		if err != nil && err != io.EOF {
			color.Redf("error: 读取文件 %s 失败: %v", filepath, err)
			return "", fmt.Errorf("error: 读取文件 %s 失败: %v", filepath, err)
		}
		if n == 0 {
			break
		}
		hash.Write(buffer[:n]) // 将读取的内容写入哈希对象
	}

	// 获取哈希值
	actualChecksum := hex.EncodeToString(hash.Sum(nil))

	// 返回哈希值
	return actualChecksum, nil
}

// getSize 函数用于获取指定文件的大小
// 参数 filePath 是文件的路径
// 返回值是一个 int64 类型的文件大小和一个 error 类型的错误信息
func getSize(filePath string) (int64, error) {
	// 调用 checkFileOrDir 函数检查文件或目录是否存在
	info, err := checkFileOrDir(filePath)
	if err != nil {
		return 0, err
	}

	// 使用 info.IsDir() 方法判断文件是否是目录
	if info.IsDir() {
		// 如果是目录，返回大小为 0 和 nil 错误
		return 0, nil
	}
	// 如果不是目录，返回文件的大小和 nil 错误
	return info.Size(), nil
}

// humanReadableSize 将字节数转换为人类可读的文件大小字符串
func humanReadableSize(size int64) string {
	// 定义单位为1024字节
	const unit = 1024
	// 如果文件大小小于1024字节，直接返回以字节为单位的字符串
	if size < unit {
		return fmt.Sprintf("%dB", size)
	}
	// 初始化除数和单位字符串为1024和"KB"
	div, unitStr := unit, "KB"
	// 如果文件大小大于等于1024KB，则转换为MB
	if size/unit >= unit {
		div, unitStr = unit*unit, "MB"
	}
	// 如果文件大小大于等于1024MB，则转换为GB
	if size/unit/unit >= unit {
		div, unitStr = unit*unit*unit, "GB"
	}
	// 返回格式化后的字符串，保留两位小数
	return fmt.Sprintf("%.2f%s", float64(size)/float64(div), unitStr)
}

// expandFilePaths 扩展文件路径，支持通配符和正则表达式
// 参数 filePaths 是需要扩展的文件路径切片
// 返回值是扩展后的文件路径切片和可能出现的错误
func expandFilePaths(filePaths []string) ([]string, error) {
	// 初始化一个空切片，用于存储扩展后的文件路径
	var expandedPaths []string

	// 遍历输入的文件路径切片
	for _, pattern := range filePaths {
		// 路径规范化解析
		pattern = filepath.Clean(pattern)

		// 尝试使用 filepath.Glob 函数解析通配符路径
		matchedPaths, err := filepath.Glob(pattern)
		if err != nil {
			// 如果解析过程中出现错误，返回错误信息
			return nil, fmt.Errorf("error: 无法解析路径模式 %s: %v", pattern, err)
		}

		// 如果匹配到路径，将所有匹配到的路径添加到扩展后的路径切片中
		if len(matchedPaths) > 0 {
			expandedPaths = append(expandedPaths, matchedPaths...)
			continue
		}

		// 如果没有匹配到路径，尝试检查是否是普通文件或目录
		if _, err := os.Stat(pattern); err == nil {
			// 如果文件或目录存在，直接将其添加到扩展后的路径切片中
			expandedPaths = append(expandedPaths, pattern)
		} else if os.IsNotExist(err) {
			// 如果路径不存在，但不是通配符模式，可以选择跳过或报错
			color.Yellowf("Warn: 路径 %s 不存在，已跳过", pattern)
			continue
		} else {
			// 如果是其他错误，返回错误信息
			return nil, fmt.Errorf("error: 检查路径 %s 时出错: %v", pattern, err)
		}
	}
	// 去重处理
	expandedPaths = uniqueStrings(expandedPaths)

	// 返回扩展后的文件路径切片和 nil 错误
	return expandedPaths, nil
}

// calculateSize 主函数，计算文件或目录的大小
func calculateSize(paths []string, jobs int) error {
	// 扩展路径，支持通配符
	expandedPaths, err := expandFilePaths(paths)
	if err != nil {
		return err
	}

	// 创建通道用于收集错误
	errChan := make(chan error)

	// 创建 WaitGroup 管理协程
	var wg sync.WaitGroup

	// 创建任务通道，限制并发数
	taskChan := make(chan string, jobs)

	// 启动工作协程
	for i := 0; i < jobs; i++ {
		wg.Add(1)
		go worker(taskChan, errChan, &wg)
	}

	// 分发任务
	for _, path := range expandedPaths {
		taskChan <- path
	}
	close(taskChan)

	// 等待所有任务完成
	go func() {
		wg.Wait()
		close(errChan)
	}()

	// 收集错误并打印
	return collectErrors(errChan)
}

// worker 协程工作函数，处理单个路径的大小计算
func worker(taskChan <-chan string, errChan chan<- error, wg *sync.WaitGroup) {
	defer wg.Done()
	for path := range taskChan {
		var size int64
		var err error

		// 调用 checkFileOrDir 函数检查文件或目录是否存在
		info, err := checkFileOrDir(path)
		if err != nil {
			errChan <- err
			continue
		}

		if info.IsDir() {
			// 如果是目录，使用 filepath.WalkDir 计算目录大小
			size, err = calculateDirSize(path)
			if err != nil {
				errChan <- fmt.Errorf("计算目录 %s 大小时出错: %v", path, err)
				continue
			}
			fmt.Printf("%s\t\t%s\n", humanReadableSize(size), path)
		} else {
			// 如果是文件，直接获取文件大小
			size = info.Size()
			fmt.Printf("%s\t\t%s\n", humanReadableSize(size), path)
		}
	}
}

// calculateDirSize 使用 filepath.WalkDir 计算目录的总大小（包括所有子目录和文件）
func calculateDirSize(dir string) (int64, error) {
	var totalSize int64

	// 规范化路径
	dir = filepath.Clean(dir)

	err := filepath.WalkDir(dir, func(path string, d os.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if !d.IsDir() {
			info, err := d.Info()
			if err != nil {
				return err
			}
			totalSize += info.Size()
		}
		return nil
	})

	return totalSize, err
}

// collectErrors 收集所有任务的错误并打印
func collectErrors(errChan <-chan error) error {
	for err := range errChan {
		color.Redf("error: %v", err)
	}
	return nil
}

// 去重函数
func uniqueStrings(strSlice []string) []string {
	keys := make(map[string]bool)
	list := []string{}
	for _, entry := range strSlice {
		if _, value := keys[entry]; !value {
			keys[entry] = true
			list = append(list, entry)
		}
	}
	return list
}
